{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generation of ideal FCC lattice sites for close-packed surface slab model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "import numpy as np\n",
    "\n",
    "print('This script generates a list of all possible ideal lattice sites layer-by-layer for a given unit cell of a FCC close-packed surface slab model.')\n",
    "print('This script is interactive and requires user input. Please read carefully before proceeding: \\n')\n",
    "\n",
    "print('Note 1: Intended to be used as an input for LAMMPS NVT trajectory clamping (via lmpclamp.py).')\n",
    "print('Note 2: Requires a LAMMPS DATA file containing equilibrated box dimensions and relaxed atomic positions.')\n",
    "print('Note 3: Assumes bottommost exposed atoms as fixed.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract unit cell information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwd = os.getcwd()\n",
    "\n",
    "# Load LAMMPS data file after NPT equilibration + ionic relaxation\n",
    "name = input('Enter the name of the LAMMPS DATA file containing the equilibrated box dimensions & relaxed atomic positions : ')\n",
    "file_eq = pwd + '/' + name\n",
    "log = open(file_eq,'r')\n",
    "log_lines = log.readlines()\n",
    "log_lines = [line.split() for line in log_lines]\n",
    "\n",
    "for line_index, line in enumerate(log_lines):\n",
    "    if line:\n",
    "        for l in line:\n",
    "            \n",
    "            if l == 'atoms':\n",
    "                N_atom = int(line[0])\n",
    "            \n",
    "            elif l == 'xlo':\n",
    "                xlo = float(line[0])\n",
    "                xhi = float(line[1])\n",
    "                Lx = xhi-xlo\n",
    "            \n",
    "            elif l == 'ylo':\n",
    "                ylo = float(line[0])\n",
    "                yhi = float(line[1])\n",
    "                Ly = yhi-ylo\n",
    "            \n",
    "            elif l == 'zlo':\n",
    "                zlo = float(line[0])\n",
    "                zhi = float(line[1])\n",
    "                Lz = zhi-zlo\n",
    "            \n",
    "            elif l == 'xy':\n",
    "                xy = float(line[0])\n",
    "                xz = float(line[1])\n",
    "                yz = float(line[2])\n",
    "            \n",
    "            elif l == 'Atoms':\n",
    "                head = line_index + 1\n",
    "                start = line_index + 2\n",
    "            \n",
    "            elif l == 'Velocities':\n",
    "                end = line_index - 1\n",
    "\n",
    "lat = np.array([[Lx, 0.0, 0.0], [xy, Ly, 0.0], [xz, yz, Lz]])    # Lattice matrix\n",
    "\n",
    "xyz = []    # N*3 matrix\n",
    "for line in log_lines[start:end]:\n",
    "\n",
    "    line[:2] = np.array(line[:2]).astype(int)    # Atom ID, atom type\n",
    "    line[2:5] = np.array(line[2:5]).astype(float)    # x, y, z\n",
    "    line[5:8] = np.array(line[5:8]).astype(int)    # ix, iy, iz (inverse sign)\n",
    "\n",
    "    x = line[2] + line[5]*Lx + line[6]*xy    # Wrapped coordinates\n",
    "    y = line[3] + line[6]*Ly\n",
    "    z = line[4] + line[7]*Lz\n",
    "    \n",
    "    xyz.append([x,y,z])\n",
    "xyz = np.array(xyz)\n",
    "\n",
    "ID2 = int(input('Enter the ID of a nearest-neighbor atom that lies in the same layer as atom #1 : '))\n",
    "ID3 = int(input('Enter the ID of another nearest-neighbor atom that lies in the same layer as atom #1 : '))\n",
    "ID4 = int(input('Enter the ID of an atom that lies one layer above atom #1 : '))\n",
    "\n",
    "for index, coord in enumerate(xyz):\n",
    "    \n",
    "    ID = log_lines[start+index][0]\n",
    "    \n",
    "    if ID == 1:\n",
    "        atom1 = coord    # atom1 = Reference atom\n",
    "    \n",
    "    elif ID == ID2:\n",
    "        atom2 = coord    # atom2 = 1st coplanar atom\n",
    "\n",
    "    elif ID == ID3:\n",
    "        atom3 = coord    # atom3 = 2nd coplanar atom\n",
    "        \n",
    "    elif ID == ID4:\n",
    "        atom4 = coord    # atom4 = Atom one layer above\n",
    "\n",
    "r12 = np.linalg.norm(np.subtract(atom1,atom2))\n",
    "r13 = np.linalg.norm(np.subtract(atom1,atom3))\n",
    "dr = np.mean([r12, r13])    # In-plane NN distance\n",
    "dr_ort = 0.5*math.sqrt(3)*dr    # Orthogonal NN distance\n",
    "\n",
    "a1 = np.subtract(atom2,atom1)    # 1st intralayer vector\n",
    "a1 = a1 / np.linalg.norm(a1)\n",
    "\n",
    "a2 = np.subtract(atom3,atom1)    # 2nd intralayer vector\n",
    "a2 = a2 / np.linalg.norm(a2)\n",
    "\n",
    "a3 = np.cross(a1,a2)    # Normal vector\n",
    "a3 = a3 / np.linalg.norm(a3)\n",
    "\n",
    "if a3[2] < 0:\n",
    "    a3 = -a3    # Ensure normal vector along +z\n",
    "    a3_flip = True\n",
    "    a1 = np.cross(a2,a3)    # Orthogonalize intralayer vector along x\n",
    "    a1 = a1 / np.linalg.norm(a1)\n",
    "else:\n",
    "    a3_flip = False\n",
    "    a2 = np.cross(a3,a1)\n",
    "    a2 = a2 / np.linalg.norm(a2)\n",
    "\n",
    "dn = np.dot(atom4,a3) - np.dot(atom1,a3)    # Interlayer distance; n = normal component\n",
    "\n",
    "log.close()\n",
    "\n",
    "print('Unit cell information & normal direction extracted. \\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rotate & generate reference sites"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if a3_flip:\n",
    "    R = np.linalg.inv(np.array([a2, a1, a3]))    # Rotation matrix\n",
    "else:\n",
    "    R = np.linalg.inv(np.array([a1, a2, a3]))\n",
    "\n",
    "xyz_rot = np.dot(xyz,R)    # Rotate unit cell\n",
    "xyz_rot = xyz_rot[np.argsort(xyz_rot[:,2])]    # Sort by z\n",
    "\n",
    "#file_rot = pwd + '/rot.xyz'\n",
    "#rot = open(file_rot,'w')\n",
    "#header_rot = '%i\\n\\n'%(N_atom)\n",
    "#rot.write(header_rot)\n",
    "#\n",
    "#for atom in range(0,N_atom):\n",
    "#    x = xyz_rot[atom][0]\n",
    "#    y = xyz_rot[atom][1]\n",
    "#    z = xyz_rot[atom][2]\n",
    "#    line_rot = '%s\\t%f\\t%f\\t%f\\n'%('Ag', x, y, z)\n",
    "#    rot.write(line_rot)\n",
    "#\n",
    "#rot.close()\n",
    "\n",
    "N_layer = 1    # Number of layers\n",
    "for atom in range(1,N_atom):\n",
    "    if xyz_rot[atom][2] - xyz_rot[atom-1][2] > 0.5*dn:\n",
    "        N_layer += 1\n",
    "\n",
    "A = []    # Site A\n",
    "for atom in range(0,N_atom):\n",
    "    if xyz_rot[atom][2] - xyz_rot[0][2] < 0.5*dn:    # Bottommost layer fixed\n",
    "        x = xyz_rot[atom][0]\n",
    "        y = xyz_rot[atom][1]\n",
    "        z = xyz_rot[atom][2] + dn    # Bottommost layer muted\n",
    "        A.append([x,y,z])\n",
    "A = np.array(A)\n",
    "\n",
    "for atom in range(0,len(A)):    # Add +/-xy periodic images to ensure no sites are missed during replication\n",
    "    \n",
    "    x0 = A[atom][0]\n",
    "    y0 = A[atom][1]\n",
    "    z0 = A[atom][2]\n",
    "    \n",
    "    # 8 additional images\n",
    "    A = np.append(A, [[x0+Lx, y0, z0]], axis=0)\n",
    "    A = np.append(A, [[x0+Lx+xy, y0+Ly, z0]], axis=0)\n",
    "    A = np.append(A, [[x0+xy, y0+Ly, z0]], axis=0)\n",
    "    A = np.append(A, [[x0-Lx+xy, y0+Ly, z0]], axis=0)\n",
    "    A = np.append(A, [[x0-Lx, y0, z0]], axis=0)\n",
    "    A = np.append(A, [[x0-Lx-xy, y0-Ly, z0]], axis=0)\n",
    "    A = np.append(A, [[x0-xy, y0-Ly, z0]], axis=0)\n",
    "    A = np.append(A, [[x0+Lx-xy, y0-Ly, z0]], axis=0)\n",
    "\n",
    "B = []    # Site B\n",
    "for atom in range(0,len(A)):\n",
    "    x = A[atom][0] + 0.5*dr\n",
    "    y = A[atom][1] + 0.5*dr/math.sqrt(3)\n",
    "    z = A[atom][2]\n",
    "    B.append([x, y, z])\n",
    "\n",
    "C = []    # Site C\n",
    "for atom in range(0,len(A)):\n",
    "    x = B[atom][0] + 0.5*dr\n",
    "    y = B[atom][1] + 0.5*dr/math.sqrt(3)\n",
    "    z = B[atom][2]\n",
    "    C.append([x, y, z])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_poscar = pwd + '/ref_site.vasp'\n",
    "poscar = open(file_poscar,'w')\n",
    "\n",
    "sys = 'Pd/Ag(111)'\n",
    "scale = 1.00000000000000\n",
    "\n",
    "ax = xhi - xlo\n",
    "ay = 0.0000000000000000\n",
    "az = 0.0000000000000000\n",
    "\n",
    "bx = xy\n",
    "by = yhi - ylo\n",
    "bz = 0.0000000000000000\n",
    "\n",
    "cx = xz\n",
    "cy = yz\n",
    "cz = zhi - zlo\n",
    "\n",
    "header_poscar = '%s\\n%f\\n%f %f %f\\n%f %f %f\\n%f %f %f\\n%s\\n%i\\n%s\\n'%(sys, scale, ax, ay, az, bx, by, bz, cx, cy, cz, 'Ag', N_atom, 'Cartesian')\n",
    "poscar.write(header_poscar)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Replicate layer-by-layer & rotate back"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file = pwd + '/ref_site.txt'\n",
    "out = open(file,'w')\n",
    "\n",
    "normal = '%s\\t%f\\t%f\\t%f\\n'%('Normal vector : ', a3[0], a3[1], a3[2])\n",
    "out.write(normal)\n",
    "\n",
    "header = 'x\\ty\\tz\\tLayer #\\tSite type\\n'\n",
    "out.write(header)\n",
    "\n",
    "# Rotate back to original orientation\n",
    "for L in range(0,N_layer+2):    # Include two extra ghost layers for the surface\n",
    "    \n",
    "    for atom in range(0,len(A)):\n",
    "        x = A[atom][0]\n",
    "        y = A[atom][1]\n",
    "        z = A[atom][2] + L*dn\n",
    "        coord = np.array([x, y, z])\n",
    "        cart = np.dot(coord, np.linalg.inv(R))    # Cartesian\n",
    "        direct = np.dot(cart, np.linalg.inv(lat))    # Direct\n",
    "        \n",
    "        for s_index, s in enumerate(direct):    # Prune sites lying outside unit cell\n",
    "            if not (0.0 < s < 1.0):\n",
    "                break\n",
    "            elif s_index == len(direct)-1:\n",
    "                line = '%f\\t%f\\t%f\\t%i\\t%s\\n'%(cart[0], cart[1], cart[2], L+1, 'A')\n",
    "                out.write(line)\n",
    "                line_poscar = '%f\\t%f\\t%f\\n'%(cart[0], cart[1], cart[2])\n",
    "                poscar.write(line_poscar)\n",
    "    \n",
    "    for atom in range(0,len(B)):\n",
    "        x = B[atom][0]\n",
    "        y = B[atom][1]\n",
    "        z = B[atom][2] + L*dn\n",
    "        coord = np.array([x, y, z])\n",
    "        cart = np.dot(coord, np.linalg.inv(R))\n",
    "        direct = np.dot(cart, np.linalg.inv(lat))\n",
    "\n",
    "        for s_index, s in enumerate(direct):\n",
    "            if not (0.0 < s < 1.0):\n",
    "                break\n",
    "            elif s_index == len(direct)-1:\n",
    "                line = '%f\\t%f\\t%f\\t%i\\t%s\\n'%(cart[0], cart[1], cart[2], L+1, 'B')\n",
    "                out.write(line)\n",
    "                line_poscar = '%f\\t%f\\t%f\\n'%(cart[0], cart[1], cart[2])\n",
    "                poscar.write(line_poscar)\n",
    "    \n",
    "    for atom in range(0,len(C)):\n",
    "        x = C[atom][0]\n",
    "        y = C[atom][1]\n",
    "        z = C[atom][2] + L*dn\n",
    "        coord = np.array([x, y, z])\n",
    "        cart = np.dot(coord, np.linalg.inv(R))\n",
    "        direct = np.dot(cart, np.linalg.inv(lat))\n",
    "\n",
    "        for s_index, s in enumerate(direct):\n",
    "            if not (0.0 < s < 1.0):\n",
    "                break\n",
    "            elif s_index == len(direct)-1:\n",
    "                line = '%f\\t%f\\t%f\\t%i\\t%s\\n'%(cart[0], cart[1], cart[2], L+1, 'C')\n",
    "                out.write(line)\n",
    "                line_poscar = '%f\\t%f\\t%f\\n'%(cart[0], cart[1], cart[2])\n",
    "                poscar.write(line_poscar)\n",
    "\n",
    "out.close()\n",
    "poscar.close()\n",
    "\n",
    "print('Reference sites generated > ref_site.txt\\n')\n",
    "\n",
    "print('lmpclamp.py can be used for LAMMPS NVT trajectory clamping.')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
